import uvicorn
from fastapi import FastAPI, Query, HTTPException
from fastapi.responses import StreamingResponse, JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from Server.AgentFlowOutput import run_judge_flow
from pydantic import BaseModel
from typing import Optional
import json

class CustomJSONResponse(JSONResponse):
    """自定义JSON响应，确保中文正确显示"""
    def render(self, content) -> bytes:
        return json.dumps(
            content,
            ensure_ascii=False,
            allow_nan=False,
            indent=None,
            separators=(',', ':'),
        ).encode('utf-8')

app = FastAPI(
    title="CTF Judge Agent API",
    description="API for CTF Performance Evaluation Agent",
    version="1.0.0",
    default_response_class=CustomJSONResponse  # 使用自定义JSON响应
)

# 配置CORS中间件，允许来自前端的请求
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 在生产环境中应该限制具体的域名
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

class JudgeRequest(BaseModel):
    query: str
    max_iterations: Optional[int] = 5

@app.get("/")
async def root():
    """根路径，返回API信息"""
    return {
        "message": "CTF Judge Agent API",
        "version": "1.0.0",
        "endpoints": {
            "/evaluate": "GET - 评估用户或题目性能",
            "/evaluate_stream": "POST - 流式评估用户或题目性能",
            "/health": "GET - 健康检查",
            "/tools": "GET - 获取可用工具列表"
        }
    }

@app.get("/health")
async def health_check():
    """健康检查端点"""
    return {"status": "healthy", "service": "CTF Judge Agent API"}

@app.get("/evaluate")
async def evaluate(
    query: str = Query(..., description="评估查询，例如：'评估用户user_001的表现'"),
    max_iterations: int = Query(5, description="最大迭代次数")
):
    """
    评估用户或题目性能的端点
    
    参数:
    - query: 评估查询字符串
    - max_iterations: 最大迭代次数，默认为5
    
    返回:
    - 流式响应，实时返回评估结果
    """
    if not query.strip():
        raise HTTPException(status_code=400, detail="Query parameter cannot be empty")
    
    if max_iterations < 1 or max_iterations > 10:
        raise HTTPException(status_code=400, detail="max_iterations must be between 1 and 10")

    return StreamingResponse(
        run_judge_flow(query=query, max_iterations=max_iterations),
        media_type="text/event-stream; charset=utf-8",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "Access-Control-Allow-Origin": "*",
            "Content-Type": "text/event-stream; charset=utf-8",
        }
    )

@app.post("/evaluate_stream")
async def evaluate_stream(request: JudgeRequest):
    """
    POST方式的流式评估端点
    
    请求体:
    {
        "query": "评估用户user_001的表现",
        "max_iterations": 5
    }
    """
    if not request.query.strip():
        raise HTTPException(status_code=400, detail="Query cannot be empty")
    
    if request.max_iterations < 1 or request.max_iterations > 10:
        raise HTTPException(status_code=400, detail="max_iterations must be between 1 and 10")

    return StreamingResponse(
        run_judge_flow(query=request.query, max_iterations=request.max_iterations),
        media_type="text/event-stream; charset=utf-8",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "Access-Control-Allow-Origin": "*",
            "Content-Type": "text/event-stream; charset=utf-8",
        }
    )

@app.get("/tools")
async def get_available_tools():
    """获取可用工具列表"""
    try:
        from Agent.AgentTools import get_tools
        tools = get_tools()
        
        tool_info = []
        for tool in tools:
            tool_info.append({
                "name": tool.__name__,
                "description": getattr(tool, 'description', 'No description available'),
                "input_schema": tool.input_schema.__name__ if hasattr(tool, 'input_schema') else 'N/A'
            })
        
        return {
            "available_tools": tool_info,
            "total_tools": len(tool_info)
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Failed to get tools: {str(e)}")

if __name__ == "__main__":
    uvicorn.run(
        app, 
        host="0.0.0.0", 
        port=8000,
        reload=True,  # 开发模式，自动重载
        log_level="info"
    )